import colorsys
import os
import time

import cv2
import numpy as np
import torch

from pr_yolo.Models.yolo_body import YoloBody
from pr_yolo.Utils.utils import (cvtColor, get_classes, resize_image)
from pr_yolo.Utils.utils_bbox import DecodeBox
import configuration as cfg

class YOLO(object):
    def __init__(self,confidence=0.5,nms_iou=0.3,letterbox_image=False):
        '''
        confidence:只有得分大于置信度的预测框会被保留下来
        nms_iou:非极大抑制所用到的nms_iou大小，越小表示越严格
        letterbox_image:该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize，在多次测试后，发现关闭letterbox_image直接resize的效果更好
        '''
        self.cuda = cfg.cuda
        self.backbone = cfg.backbone
        self.input_shape = cfg.input_img_size
        self.anchors_mask = cfg.anchors_mask
        self.anchors = cfg.yolo_anchors
        self.num_anchors = cfg.num_anchors
        self.model_path = cfg.test_model_path
        self.classes_path = cfg.classes_path
        self.confidence = confidence
        self.nms_iou = nms_iou
        self.letterbox_image = letterbox_image      #是否对图像进行不失真的resize

        #获得类别
        self.class_names, self.num_classes  = get_classes(self.classes_path)
        #解析输出的三个张量的类
        self.bbox_util                      = DecodeBox(self.anchors, self.num_classes, self.input_shape, self.anchors_mask)

        #画框设置不同的颜色
        hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
        self.generate()

    def generate(self):
        '''
        加载模型权重到结构中
        '''
        self.net    = YoloBody(self.anchors_mask, self.num_classes, self.backbone)  #声明模型

        if cfg.cuda:
            self.net = self.net.cuda()
        #==============加载模型==================
        device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model_dict = torch.load(self.model_path,map_location=device)
        self.net.load_state_dict(model_dict['net'])
        self.net    = self.net.eval()
        print('预测加载模型，来自{}'.format(self.model_path))

    def detect_image(self, image):
        '''
        检测图片，返回检测到的resize后的图片，多个对象的类别id，置信度，box坐标
        image:是一个PIL格式的图片对象
        '''
        image_shape = np.array(np.shape(image)[0:2])
        image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)#转换为RGB
        #   给图像增加灰条，实现不失真的resize,也可以直接resize进行识别
        image_data  = resize_image(image, self.input_shape, self.letterbox_image)
        #   添加上batch_size维度
        image_data  = np.expand_dims(np.transpose((np.array(image_data, dtype='float32'))/255.0, (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
            #   将图像输入网络当中进行预测
            outputs = self.net(images)  #输入是(1,3,608,608),已被归一化，输出是三个张量
            outputs = self.bbox_util.decode_box(outputs)
            #将预测框进行堆叠，然后进行非极大抑制,返回一个{ndarry:(8,7)}，8是对象数量，7是对应着box的信息
            results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
                        image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
                                                    
            if results[0] is None:
                return image,None,None,None

            top_label   = np.array(results[0][:, 6], dtype = 'int32')   #类别
            top_conf    = results[0][:, 4] * results[0][:, 5]   #置信度
            top_boxes   = results[0][:, :4]     #置信度最高的boxes它们的坐标

        return image,top_label,top_conf,top_boxes

    def visualize(self,image,labels,confidence,boxed_cord):
        '''
        根据检测生成的box信息，可视化这个结果
        image:用于检测的图像{ndarray:(1080,1920,3)}是RGB格式
        labels:标签id,{ndarray:(7,)} [2 2 2 2 2 2 2]
        confidence:标签置信度,{ndarray:(7,)} [0.99 0.81 ... 0.56]
        boxes_cord:标签坐标，{ndarray:(7,4)} 为x_min,y_min,x_max,y_max
        -----------------
        return:
            返回生成的图片，可决定是否保存
        '''
        #图像绘制
        for i, c in list(enumerate(labels)):
            predicted_class = self.class_names[int(c)]
            box = boxed_cord[i]
            score = confidence[i]

            top, left, bottom, right = box

            top = max(0, np.floor(top).astype('int32'))
            left = max(0, np.floor(left).astype('int32'))
            bottom = min(image.shape[0], np.floor(bottom).astype('int32'))
            right = min(image.shape[1], np.floor(right).astype('int32'))
            #开始绘画，我们将标签信息放在框框的里面，这样省略了很多代码
            label = '{} {:.2f}'.format(predicted_class, score)  #显示的文字信息
            cv2.rectangle(image,(left,top),(right,bottom),self.colors[c],3)
            cv2.putText(image,label,(left,top),
                              cv2.FONT_HERSHEY_PLAIN,1,color=self.colors[c],thickness=1)
        return image

    def get_FPS(self, image, test_interval):
        image_shape = np.array(np.shape(image)[0:2])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data  = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose((np.array(image_data, dtype='float32'))/255.0, (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
            #---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            #---------------------------------------------------------#
            outputs = self.net(images)
            outputs = self.bbox_util.decode_box(outputs)
            #---------------------------------------------------------#
            #   将预测框进行堆叠，然后进行非极大抑制
            #---------------------------------------------------------#
            results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, 
                        image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
                                                    
        t1 = time.time()
        for _ in range(test_interval):
            with torch.no_grad():
                #---------------------------------------------------------#
                #   将图像输入网络当中进行预测！
                #---------------------------------------------------------#
                outputs = self.net(images)
                outputs = self.bbox_util.decode_box(outputs)
                #---------------------------------------------------------#
                #   将预测框进行堆叠，然后进行非极大抑制
                #---------------------------------------------------------#
                results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, 
                            image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
                            
        t2 = time.time()
        tact_time = (t2 - t1) / test_interval
        return tact_time

    def get_map_txt(self, image_id, image, class_names, map_out_path):
        f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w") 
        image_shape = np.array(np.shape(image)[0:2])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data  = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose((np.array(image_data, dtype='float32'))/255.0, (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
            #---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            #---------------------------------------------------------#
            outputs = self.net(images)
            outputs = self.bbox_util.decode_box(outputs)
            #---------------------------------------------------------#
            #   将预测框进行堆叠，然后进行非极大抑制
            #---------------------------------------------------------#
            results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, 
                        image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
                                                    
            if results[0] is None: 
                return 

            top_label   = np.array(results[0][:, 6], dtype = 'int32')
            top_conf    = results[0][:, 4] * results[0][:, 5]
            top_boxes   = results[0][:, :4]

        for i, c in list(enumerate(top_label)):
            predicted_class = self.class_names[int(c)]
            box             = top_boxes[i]
            score           = str(top_conf[i])

            top, left, bottom, right = box
            if predicted_class not in class_names:
                continue

            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))

        f.close()
        return 
